// enum-modules-win.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"

// To ensure correct resolution of symbols, add Psapi.lib to TARGETLIBS
// and compile with -DPSAPI_VERSION=1
#pragma comment(lib, "psapi") 


int PrintModules( DWORD processID )
{
    HMODULE hMods[1024];
    HANDLE hProcess;
    DWORD cbNeeded;
    unsigned int i;
	TCHAR szProcName[MAX_PATH];
	DWORD szProcNameSize = MAX_PATH;

    // Get a handle to the process.

    hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
                            PROCESS_VM_READ,
                            FALSE, processID );
    if (NULL == hProcess)
        return 1;

	// Print the process identifier.
    printf( "Process ID: %u", processID,szProcName );

   // get the name of the process
	if ( QueryFullProcessImageName( hProcess, 0x0, szProcName, &szProcNameSize))
		{
			// Print the process name
			_tprintf( TEXT(", Process Name: %s"),szProcName );
		}
	printf("\n");

   // Get a list of all the modules in this process.

    if( EnumProcessModules(hProcess, hMods, sizeof(hMods), &cbNeeded))
    {
        for ( i = 0; i < (cbNeeded / sizeof(HMODULE)); i++ )
        {
            TCHAR szModName[MAX_PATH];

            // Get the full path to the module's file.

            if ( GetModuleFileNameEx( hProcess, hMods[i], szModName,
                                      sizeof(szModName) / sizeof(TCHAR)))
            {
                // Print the module name and handle value.

                _tprintf( TEXT("\t%s (0x%08X)\n"), szModName, hMods[i] );
            }
        }
    }
    
    // Release the handle to the process.

    CloseHandle( hProcess );

    return 0;
}

bool IsModuleLoaded(DWORD processID, std::wstring searchModule)
{
	HMODULE hMods[1024];
    HANDLE hProcess;
    DWORD cbNeeded;
    unsigned int i;
	std::wstring isModule;
	bool bFound =false;

    // Get a handle to the process.

    hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
                            PROCESS_VM_READ,
                            FALSE, processID );
    if (NULL == hProcess)
        return 1;


	TCHAR szProcName[MAX_PATH];
	DWORD szProcNameSize = MAX_PATH;
	// get the name of the process
	// Print the process identifier.
    
	if ( QueryFullProcessImageName( hProcess, 0x0, szProcName, &szProcNameSize))
		{
			// Print the process name
#if _DEBUG
			printf( "Process ID: %u", processID,szProcName );
			_tprintf( TEXT(", Process Name: %s"),szProcName );
			printf("\n");
#endif
		}


   // Get a list of all the modules in this process.

    if( EnumProcessModules(hProcess, hMods, sizeof(hMods), &cbNeeded))
    {
        for ( i = 0; i < (cbNeeded / sizeof(HMODULE)); i++ )
        {
            TCHAR szModName[MAX_PATH];
            // Get the name of the module's file.

            if ( GetModuleBaseName( hProcess, hMods[i], szModName,
                                      sizeof(szModName) / sizeof(TCHAR)))
            {
                //compare the module's basename to what we are looking for
				isModule = std::wstring(szModName);
				std::transform(isModule.begin(), isModule.end(), isModule.begin(), ::tolower); // to lowercase
				if (!searchModule.compare(isModule))
				{
					bFound=true;
					break;
				}
            }
        }
    }
    
    // Release the handle to the process.
	CloseHandle( hProcess );

	std::wcout << TEXT("Found the module '") << searchModule.c_str() << TEXT("' for process '") << szProcName << TEXT("'?:");
	if (bFound) {std::wcout << TEXT("YES\n");} else {std::wcout << TEXT("NO\n");}

    return !bFound;
}

void PrintHelp()
{
	std::cerr <<"enum-modules-win.exe Version 1.1\n";
	std::cerr <<"This tool queries the loaded modules of running processes\n";
	std::cerr <<"Usage:\n";
	std::cerr <<"enum-modules-win.exe --list          : Lists all running processes and their loaded modules\n";
	std::cerr <<"enum-modules-win.exe --pid <PID>     : Lists all loaded modules of the process with the pid specified in <PID> (e.g. 1234)\n";
	std::cerr <<"enum-modules-win.exe --name <NAME>   : Lists all loaded modules of the process with the name specified in <NAME> (e.g. explorer.exe)\n";
	std::cerr <<"enum-modules-win.exe --pid <PID> <MODNAME>   : Checks if the process with pid <PID> has module with name <MODNAME> loaded\n";
	std::cerr <<"                                               Returns with exit code 0 when loaded and exit code 1 if not loaded\n";
	std::cerr <<"enum-modules-win.exe --name <NAME> <MODNAME> : Checks if the process with name <NAME> has module with name <MODNAME> loaded\n";
	std::cerr <<"                                               Returns with exit code 0 when loaded and exit code 1 if not loaded\n";
	std::cerr <<"Example:\n";
	std::cerr <<"enum-modules-win.exe --name explorer.exe duser.dll   : This call will check if the process explorer.exe has a module with name duser.dll loaded.\n";
	std::cerr <<"                                                       If not it will exit with error code 1, if the module is loaded it will exit with 0\n";
	std::cerr <<"Notes:\n";
	std::cerr <<"Use the Windows environment variable %ERRORLEVEL% immediatelly after this program finishes execution to check for the exit code in batch scripts";
	std::cerr <<"\n";
}

int ListAll()
{
	DWORD aProcesses[1024]; 
    DWORD cbNeeded; 
    DWORD cProcesses;
    unsigned int i;
	// Get the list of process identifiers.
    if ( !EnumProcesses( aProcesses, sizeof(aProcesses), &cbNeeded ) )
        return 1;

    // Calculate how many process identifiers were returned.
	cProcesses = cbNeeded / sizeof(DWORD);

    // Print the names of the modules for each process.
	for ( i = 0; i < cProcesses; i++ )
    {
        PrintModules( aProcesses[i] );
    }
	return 0;
}

int GetPidByName(std::wstring name)
{
    DWORD pid = 0;
	std::wstring procname;

    // Create toolhelp snapshot.
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    PROCESSENTRY32 process;
    ZeroMemory(&process, sizeof(process));
    process.dwSize = sizeof(process);

    // Walkthrough all processes.
    if (Process32First(snapshot, &process))
    {
        do
        {
            // Compare process.szExeFile based on format of name, i.e., trim file path
            // trim .exe if necessary, etc.
			procname = std::wstring(process.szExeFile);
			std::transform(procname.begin(), procname.end(), procname.begin(), ::tolower); // to lowercase
			if (!name.compare(procname))
            {
               pid = process.th32ProcessID;
               break;
            }
        } while (Process32Next(snapshot, &process));
    }

    CloseHandle(snapshot);

    if (pid != 0)
    {
         return pid;
    }

    // Not found
	return NULL;
}


int _tmain(int argc, WCHAR* argv[])
{
	std::wstring		procName;
	std::wstring		modName;
	std::wstring		argFlag;
	bool			bPid;
	bool			bModSearch;

	//handle command line input
	if (argc < 2) {
		PrintHelp();
		return 99;
	}

	argFlag = std::wstring(argv[1]);
	if (!argFlag.compare(TEXT("--pid")))
	{
		bPid = true;
	} else if (!argFlag.compare(TEXT("--name")))
	{ 
		bPid = false;
	} else if (!argFlag.compare(TEXT("--list")))
	{ 
		return ListAll();  // just list all of them
	} else
	{
		PrintHelp();
		return 99;
	}

	if (argc < 3){
		PrintHelp();
		return 99;
	}
	
	// second argument
	procName = std::wstring(argv[2]);

	// third, optional argument
	bModSearch= false;
	if (argc == 4)
	{
		modName = std::wstring(argv[3]);
		bModSearch= true;
	}

#if _DEBUG
	if (bPid) std::cout <<"running in PID mode\n";
	if (!bPid) std::cout <<"running in NAME mode\n";
	std::wcout << "Procname/Pid:" << procName.c_str() << "\n";
	if (bModSearch) std::wcout << "Searching for module with name : " << modName.c_str() << "\n";
#endif

	if (bPid){
		// only list this PID
		if (bModSearch)
		{
			return IsModuleLoaded(_wtoi(procName.c_str()),modName);
		} 
		else
		{
			if (PrintModules(_wtoi(procName.c_str())) != 0)
			{
				std::wcerr << "Could not find process with PID " << procName.c_str() << "\n";
				return 1;
			}
		}
	} else 
	{
		// search process by name
		std::transform(procName.begin(), procName.end(), procName.begin(), ::tolower); // to lowercase
		int pid = GetPidByName(procName);

		if (bModSearch)
		{
			return IsModuleLoaded(pid,modName);
		} 
		else
		{
			if (PrintModules(pid) != 0)
			{
				std::wcerr << "Could not find process with PID " << procName.c_str() << "\n";
				return 1;
			}
		}
	}

    return 0;
}

